import math
from collections import OrderedDict

import torch
import torch.nn as nn
import torchvision.transforms as tt

#  proposed fire block
class FireBlock(nn.Module):
    def __init__(self, inplanes, planes):    # self.inplanes=64，记住在下面修改inplanes, planes
        super(FireBlock, self).__init__()      
        # -----------------------------------------#
        # input node: 64*55*55
        # 定义了两组卷积，以及concat：
        # 第一组：16,1,1
        # 第二组：64,1,1 + RELU 并列  64,3,3 + RELU 
        # 最后：concat+BN
        # output node: 128*55*55
        # -----------------------------------------#
        # 第一组
        # conv--16*1*1 :  64*55*55 -> 16*55*55
        self.conv1 = nn.Conv2d(inplanes, planes[0], kernel_size=1,
                               stride=1, padding=0, bias=False)
        # self.bn1 = nn.BatchNorm2d(planes[0])
        self.relu1 = nn.ReLU(inplace=True)   # inplace=True表示覆盖计算结果，节省内存
        
        # 第二组：
        # conv--64*1*1 ： 16*55*55 -> 64*55*55
        self.conv21 = nn.Conv2d(planes[0], planes[1], kernel_size=1,
                               stride=1, padding=0, bias=False)
        # self.bn21 = nn.BatchNorm2d(planes[1])
        self.relu21 = nn.ReLU(inplace=True)

        # conv--64*3*3 ： 16*55*55 -> 64*55*55
        self.conv22 = nn.Conv2d(planes[0], planes[1], kernel_size=3,
                               stride=1, padding=1, bias=False)
        # self.bn22 = nn.BatchNorm2d(planes[1])
        self.relu22 = nn.ReLU(inplace=True)

        # BN
        self.bn2 = nn.BatchNorm2d(2 * planes[1])

    # fire block前向传播
    def forward(self, x):
        # 第一个卷积层       
        out = self.conv1(x)
        # out = self.bn1(out)
        out = self.relu1(out)
        
        # 第二个并行卷积层
        # 第一个卷积
        out1 = self.conv21(out)
        # out1 = self.bn2(out1)
        out1 = self.relu21(out1)
        # 第二个卷积
        out2 = self.conv22(out)
        # out2 = self.bn2(out2)
        out2 = self.relu22(out2)
        
        # concat操作
        out = torch.cat([out1, out2], 1)
        # BN
        out_bn2 = self.bn2(out)
        
        # 转化数据为tensor类型
        # out_bn2 = tf(out_bn2)
        return out_bn2
        # FireBlock构建成功

# 建4个DDF模型
class DDF1(nn.Module):           
    def __init__(self, delta):
        # nn.Module的子类函数必须在构造函数中执行父类的构造函数
        #  last_plane代表输入DDF模块的最后一个3*3卷积的输入通道
        super(DDF1, self).__init__()  # 等价于nn.Module.__init__()
        self.inplanes = [64, 192]  #输入每个fireblock的特征图通道数
        self.planes = [16, 64]   
        

        #  一个DDF模块有2个fireblock：fireblock1和fireblock2
        #  定义fireblock 1~2,对于其他DDF模块，直接重新赋self.inplanes[2]和[3],
        # 再分别赋给self.fireblock[0]和[1]，以此类推
        self.fireblock1 = FireBlock(self.inplanes[0], self.planes)
        self.fireblock2 = FireBlock(self.inplanes[1], self.planes)
        
        # 定义两个fireblock的按位元素相加
        # self.fb_add = torch.add(self.fireblock1, self.fireblock2)

        # 定义conv: (0.5*d)*3*3 + RELU
        # 448? 按照原理来说应该是576， 576/2=288
        self.conv = nn.Conv2d(448, 224, kernel_size=3, stride=1, padding=0, dilation=delta, bias=False)
        # self.conv = nn.Conv2d(576, 288, kernel_size=3, stride=1, padding=0, dilation=delta, bias=False)
        self.relu = nn.ReLU(inplace=True)
        
        # 定义DDF中的两个BN
        self.bn_1 = nn.BatchNorm2d(self.inplanes[0]+128)
        self.bn_2 = nn.BatchNorm2d(self.inplanes[1]+128*2)

        # 定义Dropout(是对神经元输出值的激活值进行失活处理),dropout前后shape不变
        self.dropout = nn.Dropout(p=0.5)
                
        # 进行权值初始化 
        # 从self.modules()中遍历每一层
        # 然后判断每层属于什么类型，是否是Conv2d，是否是BatchNorm2d，
        # 然后根据不同类型的层，设定不同的权值初始化方法
        for m in self.modules():
            if isinstance(m, nn.Conv2d):     # isinstance用于判断m是否是nn.Conv2d类型的数据
                # 用于计算权值所属的正态分布的标准差
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)    # 全填充1
                m.bias.data.zero_()
                # _表示在一个tensor上操作了之后，直接修改这个tensor

   
    # 只要在nn.Module的子类中定义了forward函数，backward函数就会被自动实现(利用Autograd)
    def forward(self, x):
        x1 = self.fireblock1(x)

        x = torch.cat([x1, x], 1)
        x = self.bn_1(x)

        x2 = self.fireblock2(x)
        
        x_add = torch.add(x1, x2)
        x = torch.cat([x2, x_add, x], 1)
        x = self.bn_2(x)

        x = self.conv(x)
        x = self.relu(x)

        ddf1_out = self.dropout(x)

        return ddf1_out

class DDF2(nn.Module):           
    def __init__(self, delta):
        # nn.Module的子类函数必须在构造函数中执行父类的构造函数
        #  last_plane代表输入DDF模块的最后一个3*3卷积的输入通道
        super(DDF2, self).__init__()  # 等价于nn.Module.__init__()
        self.inplanes = [224, 480]
        # self.inplanes = [288, 544]  #输入每个fireblock的特征图通道数
        self.planes = [32, 128]   
        

        #  一个DDF模块有2个fireblock：fireblock1和fireblock2
        #  定义fireblock 1~2,对于其他DDF模块，直接重新赋self.inplanes[2]和[3],
        # 再分别赋给self.fireblock[0]和[1]，以此类推
        self.fireblock1 = FireBlock(self.inplanes[0], self.planes)
        self.fireblock2 = FireBlock(self.inplanes[1], self.planes)
        
        # 定义两个fireblock的按位元素相加
        # self.fb_add = torch.add(self.fireblock1, self.fireblock2)

        # 定义conv: (0.5*d)*3*3 + RELU
        self.conv = nn.Conv2d(992, 496, kernel_size=3, stride=1, padding=0, dilation=delta, bias=False)
        # self.conv = nn.Conv2d(1312, 656, kernel_size=3, stride=1, padding=0, dilation=delta, bias=False)
        self.relu = nn.ReLU(inplace=True)
        
        # 定义DDF中的两个BN
        self.bn_1 = nn.BatchNorm2d(self.inplanes[0]+256)
        self.bn_2 = nn.BatchNorm2d(self.inplanes[1]+256*2)

        # 定义Dropout(是对神经元输出值的激活值进行失活处理),dropout前后shape不变
        self.dropout = nn.Dropout(p=0.5)
                
        # 进行权值初始化 
        # 从self.modules()中遍历每一层
        # 然后判断每层属于什么类型，是否是Conv2d，是否是BatchNorm2d，
        # 然后根据不同类型的层，设定不同的权值初始化方法
        for m in self.modules():
            if isinstance(m, nn.Conv2d):     # isinstance用于判断m是否是nn.Conv2d类型的数据
                # 用于计算权值所属的正态分布的标准差
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)    # 全填充1
                m.bias.data.zero_()
                # _表示在一个tensor上操作了之后，直接修改这个tensor

   
    # 只要在nn.Module的子类中定义了forward函数，backward函数就会被自动实现(利用Autograd)
    def forward(self, x):
        x1 = self.fireblock1(x)

        x = torch.cat([x1, x], 1)
        x = self.bn_1(x)

        x2 = self.fireblock2(x)
        
        x_add = torch.add(x1, x2)
        x = torch.cat([x2, x_add, x], 1)
        x = self.bn_2(x)

        x = self.conv(x)
        x = self.relu(x)

        ddf2_out = self.dropout(x)

        return ddf2_out

class DDF3(nn.Module):
    def __init__(self, delta):
        # nn.Module的子类函数必须在构造函数中执行父类的构造函数
        #  last_plane代表输入DDF模块的最后一个3*3卷积的输入通道
        super(DDF3, self).__init__()  # 等价于nn.Module.__init__()
        self.inplanes = [496, 880]
        # self.inplanes = [656, 1040]  #输入每个fireblock的特征图通道数
        self.planes = [48, 192]   #这个是固定的
        

        #  一个DDF模块有2个fireblock：fireblock1和fireblock2
        #  定义fireblock 1~2,对于其他DDF模块，直接重新赋self.inplanes[2]和[3],
        # 再分别赋给self.fireblock[0]和[1]，以此类推
        self.fireblock1 = FireBlock(self.inplanes[0], self.planes)
        self.fireblock2 = FireBlock(self.inplanes[1], self.planes)
        
        # 定义两个fireblock的按位元素相加
        # self.fb_add = torch.add(self.fireblock1, self.fireblock2)

        # 定义conv: (0.5*d)*3*3 + RELU
        self.conv = nn.Conv2d(1648, 824, kernel_size=3, stride=1, padding=0, dilation=delta, bias=False)
        # self.conv = nn.Conv2d(2192, 1096, kernel_size=3, stride=1, padding=0, dilation=delta, bias=False)
        self.relu = nn.ReLU(inplace=True)
        
        # 定义DDF中的两个BN
        self.bn_1 = nn.BatchNorm2d(self.inplanes[0]+384)
        self.bn_2 = nn.BatchNorm2d(self.inplanes[1]+384*2)

        # 定义Dropout(是对神经元输出值的激活值进行失活处理),dropout前后shape不变
        self.dropout = nn.Dropout(p=0.5)
                
        # 进行权值初始化 
        # 从self.modules()中遍历每一层
        # 然后判断每层属于什么类型，是否是Conv2d，是否是BatchNorm2d，
        # 然后根据不同类型的层，设定不同的权值初始化方法
        for m in self.modules():
            if isinstance(m, nn.Conv2d):     # isinstance用于判断m是否是nn.Conv2d类型的数据
                # 用于计算权值所属的正态分布的标准差
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)    # 全填充1
                m.bias.data.zero_()
                # _表示在一个tensor上操作了之后，直接修改这个tensor

   
    # 只要在nn.Module的子类中定义了forward函数，backward函数就会被自动实现(利用Autograd)
    def forward(self, x):
        x1 = self.fireblock1(x)

        x = torch.cat([x1, x], 1)
        x = self.bn_1(x)

        x2 = self.fireblock2(x)
        
        x_add = torch.add(x1, x2)
        x = torch.cat([x2, x_add, x], 1)
        x = self.bn_2(x)

        x = self.conv(x)
        x = self.relu(x)

        ddf3_out = self.dropout(x)

        return ddf3_out

class DDF4(nn.Module):
    def __init__(self, delta):
        # nn.Module的子类函数必须在构造函数中执行父类的构造函数
        #  last_plane代表输入DDF模块的最后一个3*3卷积的输入通道
        super(DDF4, self).__init__()  # 等价于nn.Module.__init__()
        self.inplanes = [824, 1336]
        # self.inplanes = [1096, 1608]  #输入每个fireblock的特征图通道数
        self.planes = [64, 256]   #这个是固定的
        

        #  一个DDF模块有2个fireblock：fireblock1和fireblock2
        #  定义fireblock 1~2,对于其他DDF模块，直接重新赋self.inplanes[2]和[3],
        # 再分别赋给self.fireblock[0]和[1]，以此类推
        self.fireblock1 = FireBlock(self.inplanes[0], self.planes)
        self.fireblock2 = FireBlock(self.inplanes[1], self.planes)
        
        # 定义两个fireblock的按位元素相加? 
        # 如果是这样的话，那么3*3conv那一步计算的通道就不对了，这一步应该是通道叠加才对？
        # self.fb_add = torch.add(self.fireblock1, self.fireblock2)

        # 定义conv: (0.5*d)*3*3 + RELU
        self.conv = nn.Conv2d(2360, 1180, kernel_size=3, stride=1, padding=0, dilation=delta, bias=False)
        # self.conv = nn.Conv2d(3144, 1572, kernel_size=3, stride=1, padding=0, dilation=delta, bias=False)
        self.relu = nn.ReLU(inplace=True)
        
        # 定义DDF中的两个BN
        self.bn_1 = nn.BatchNorm2d(self.inplanes[0]+512)
        self.bn_2 = nn.BatchNorm2d(self.inplanes[1]+512*2)

        # 定义Dropout(是对神经元输出值的激活值进行失活处理),dropout前后shape不变
        self.dropout = nn.Dropout(p=0.5)
                
        # 进行权值初始化 
        # 从self.modules()中遍历每一层
        # 然后判断每层属于什么类型，是否是Conv2d，是否是BatchNorm2d，
        # 然后根据不同类型的层，设定不同的权值初始化方法
        for m in self.modules():
            if isinstance(m, nn.Conv2d):     # isinstance用于判断m是否是nn.Conv2d类型的数据
                # 用于计算权值所属的正态分布的标准差
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)    # 全填充1
                m.bias.data.zero_()
                # _表示在一个tensor上操作了之后，直接修改这个tensor

   
    # 只要在nn.Module的子类中定义了forward函数，backward函数就会被自动实现(利用Autograd)
    def forward(self, x):
        x1 = self.fireblock1(x)

        x = torch.cat([x1, x], 1)
        x = self.bn_1(x)

        x2 = self.fireblock2(x)
        
        x_add = torch.add(x1, x2)
        x = torch.cat([x2, x_add, x], 1)
        x = self.bn_2(x)

        x = self.conv(x)
        x = self.relu(x)

        ddf4_out = self.dropout(x)

        return ddf4_out